import math
import numpy as np
import torch


def xyxy2xywh(x):
    """
    Convert bounding box format from [x1, y1, x2, y2] to [cx, cy, w, h]
    """
    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
    y[..., 0] = (x[..., 0] + x[..., 2]) / 2
    y[..., 1] = (x[..., 1] + x[..., 3]) / 2
    y[..., 2] = x[..., 2] - x[..., 0]
    y[..., 3] = x[..., 3] - x[..., 1]
    return y


def xywh2xyxy(x):
    """
    Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2]
    """
    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2
    y[..., 2] = x[..., 0] + x[..., 2] / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2
    return y


def bbox_wh_iou(wh1, wh2):
    wh2 = wh2.t()
    w1, h1 = wh1[0], wh1[1]
    w2, h2 = wh2[0], wh2[1]
    inter_area = torch.min(w1, w2) * torch.min(h1, h2)
    union_area = (w1 * h1 + 1e-16) + w2 * h2 - inter_area
    return inter_area / union_area


def bbox_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes, box1: nx4, box2: nx4
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) * \
                 torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
    # Union Area
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou


def bbox_giou(box1, box2, x1y1x2y2=True):
    """
    Returns the GIoU of two bounding boxes, box1: nx4, box2: nx4
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    union_rect_x1 = torch.min(b1_x1, b2_x1)
    union_rect_y1 = torch.min(b1_y1, b2_y2)
    union_rect_x2 = torch.max(b1_x2, b2_x2)
    union_rect_y2 = torch.max(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) * \
                 torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
    union_area = torch.clamp(union_rect_x2 - union_rect_x1, min=0) * \
                 torch.clamp(union_rect_y2 - union_rect_y1, min=0)
    # Union Area
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
    giou = iou - (union_area - b1_area - b2_area + inter_area) / (union_area + 1e-16)
    giou = torch.clamp(giou, min=-1.0, max=1.0)
    return giou


def bbox_diou(box1, box2, x1y1x2y2=True):
    """
    Returns the DIoU of two bounding boxes, box1: nx4, box2: nx4
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        b1_cx, b1_cy = box1[:, 0], box1[:, 1]
        b2_cx, b2_cy = box2[:, 0], box2[:, 1]
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
        b1_cx, b1_cy = (b1_x1 + b1_x2) * 0.5, (b1_y1 + b1_y2) * 0.5
        b2_cx, b2_cy = (b2_x1 + b2_x2) * 0.5, (b2_y1 + b2_y2) * 0.5

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    union_rect_x1 = torch.min(b1_x1, b2_x1)
    union_rect_y1 = torch.min(b1_y1, b2_y2)
    union_rect_x2 = torch.max(b1_x2, b2_x2)
    union_rect_y2 = torch.max(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) * \
                 torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
    # Center distance
    dist_center = (b1_cx - b2_cx) ** 2 + (b1_cy - b2_cy) ** 2
    # union rect distance
    dist_diag = (union_rect_x1 - union_rect_x2) ** 2 + (union_rect_y1 - union_rect_y2) ** 2
    # Union Area
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
    diou = iou - dist_center / (dist_diag + 1e-16)
    diou = torch.clamp(diou, min=-1.0, max=1.0)
    return diou


def bbox_ciou(box1, box2, x1y1x2y2=True):
    """
     Returns the DIoU of two bounding boxes, box1: nx4, box2: 1x4
     """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        b1_cx, b1_cy = box1[:, 0], box1[:, 1]
        b2_cx, b2_cy = box2[:, 0], box2[:, 1]
        b1_w, b1_h = box1[:, 2], box1[:, 3]
        b2_w, b2_h = box2[:, 2], box2[:, 3]
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
        b1_cx, b1_cy = (b1_x1 + b1_x2) * 0.5, (b1_y1 + b1_y2) * 0.5
        b2_cx, b2_cy = (b2_x1 + b2_x2) * 0.5, (b2_y1 + b2_y2) * 0.5
        b1_w, b1_h = b1_x2 - b1_x1, b1_y2 - b1_y1
        b2_w, b2_h = b2_x2 - b2_x1, b2_y2 - b2_y1

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    union_rect_x1 = torch.min(b1_x1, b2_x1)
    union_rect_y1 = torch.min(b1_y1, b2_y2)
    union_rect_x2 = torch.max(b1_x2, b2_x2)
    union_rect_y2 = torch.max(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) * \
                 torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
    # Center distance
    dist_center = (b1_cx - b2_cx) ** 2 + (b1_cy - b2_cy) ** 2
    # union rect distance
    dist_diag = (union_rect_x1 - union_rect_x2) ** 2 + (union_rect_y1 - union_rect_y2) ** 2
    # Union Area
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    v = (4.0 / (math.pi ** 2)) * (torch.atan(b2_w / b2_h) - torch.atan(b1_w / b1_h)) ** 2
    alpha = v / (1 - iou + v)

    ciou = iou - dist_center / (dist_diag + 1e-16) + alpha * v
    ciou = torch.clamp(ciou, min=-1.0, max=1.0)

    return ciou


if __name__ == "__main__":
    bbox = torch.Tensor([[60, 70, 100, 100],
                         [40, 31, 5, 6],
                         [57, 39, 2, 9]])
    gt = torch.Tensor([[75, 80, 110, 120]])

    iou = bbox_iou(bbox, gt, False)
    giou = bbox_giou(bbox, gt, False)
    diou = bbox_diou(bbox, gt, False)
    ciou = bbox_ciou(bbox, gt, False)

    pass